"""
ACTWrapper.py

加载 act om 模型并推理
"""

import numpy as np
import torch
from torch import Tensor
import OMmodel
from lerobot.policies.act.configuration_act import ACTConfig

def logger(msg: str):
    print(f'[ACTWrapper]: {msg}')

class ACTWrapper():
    def __init__(self, model_path: str, config: ACTConfig):
        self.om_model = OMmodel(model_path)
        chunk_size = config.chunk_size
        action_dim = config.output_features["action"].shape
        self.output_shape = [1, chunk_size, *action_dim]
        logger(f"Loaded ACT OM model from {model_path}, output shape: {self.output_shape}")

    def predict(self, batch: dict[str, Tensor]) -> tuple:
        input_arr = []

        key_to_exclude = [
            'observation.images', 'action', 'next.reward', 'next.done', 'next.truncated', 'info', 'task'
        ]

        for key in batch:
            # TODO: 尝试从 om model descriptor 中获取输入名称
            if key in key_to_exclude:
                continue
            
            input_arr.append(batch[key].cpu().numpy())
        
        output = self.om_model.forward(input_arr)[0]

        o_tensor = torch.from_numpy(np.array(output, dtype=np.float32)).reshape(*self.output_shape)
        return (o_tensor,)